#include "client.h"

// 定义静态成员变量
const int EchoClient::THREAD_NUM;

EchoClient::EchoClient(const char* ip, const std::vector<int>& ports, int numConnections, int maxEchoCount) 
    : ip_(ip), ports_(ports), numConnections_(numConnections), 
      epollFd_(-1), maxEchoCount_(maxEchoCount),
      charSet_("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()_+-=") {
    std::random_device rd;
    rng_.seed(rd());
}

EchoClient::~EchoClient() {
    if (epollFd_ != -1) close(epollFd_);
}

void EchoClient::setNonBlocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags != -1) fcntl(fd, F_SETFL, flags | O_NONBLOCK);
}

bool EchoClient::checkConnectSuccess(int fd) {
    int error = 0;
    socklen_t len = sizeof(error);
    return getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &len) == 0 && error == 0;
}

int EchoClient::recvAll(int fd, char* buffer, int expectedLen) {
    int total = 0;
    int retryCount = 0;
    while (total < expectedLen && retryCount < 5) {
        int n = read(fd, buffer + total, expectedLen - total);
        if (n == -1) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                usleep(10000);
                retryCount++;
                continue;
            }
            return -1;
        }
        if (n == 0) return total;
        total += n;
        retryCount = 0;
    }
    return total;
}

std::string EchoClient::generateTestData(int connId, int echoCount) {
    std::uniform_int_distribution<int> lenDist(5, 15);
    int dataLen = lenDist(rng_);
    std::string data;
    data += "conn" + std::to_string(connId) + "_echo" + std::to_string(echoCount) + "_";
    
    std::uniform_int_distribution<int> charDist(0, charSet_.size() - 1);
    for (int i = data.size(); i < dataLen; ++i) {
        data += charSet_[charDist(rng_)];
    }
    return data;
}

bool EchoClient::init() {
    epollFd_ = epoll_create1(0);
    if (epollFd_ == -1) {
        std::cerr << "Epoll create failed: " << strerror(errno) << std::endl;
        return false;
    }
    std::cout << "Epoll initialized (fd: " << epollFd_ << ")" << std::endl;
    return true;
}

int EchoClient::connectToServer(int connId) {
    std::uniform_int_distribution<int> portDist(0, ports_.size() - 1);
    int port = ports_[portDist(rng_)];

    int sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (sockfd == -1) {
        if (connId % 100 == 0) {  // 减少日志
            std::cerr << "Conn " << connId << " create socket failed" << std::endl;
        }
        return -1;
    }

    setNonBlocking(sockfd);

    // 增大发送缓冲区
    int bufSize = 8192;
    setsockopt(sockfd, SOL_SOCKET, SO_SNDBUF, &bufSize, sizeof(bufSize));

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    if (inet_pton(AF_INET, ip_.c_str(), &addr.sin_addr) <= 0) {
        close(sockfd);
        return -1;
    }

    // 非阻塞连接
    int ret = connect(sockfd, (struct sockaddr*)&addr, sizeof(addr));
    if (ret == -1 && errno != EINPROGRESS) {
        close(sockfd);
        return -1;
    }

    // 加锁保护共享数据
    {
        std::lock_guard<std::mutex> lock(connMutex_);
        connState_[sockfd] = std::make_pair("", 0);
        connCreateTime_[sockfd] = time(nullptr);
    }

    // 添加到epoll（关键：检查返回值）
    struct epoll_event ev;
    ev.events = EPOLLOUT | EPOLLIN | EPOLLERR;
    ev.data.fd = sockfd;
    if (epoll_ctl(epollFd_, EPOLL_CTL_ADD, sockfd, &ev) == -1) {
        std::cerr << "Conn " << connId << " add to epoll failed: " << strerror(errno) << std::endl;
        close(sockfd);
        std::lock_guard<std::mutex> lock(connMutex_);
        connState_.erase(sockfd);
        connCreateTime_.erase(sockfd);
        return -1;
    }

    // 每100个连接打印一次（验证连接创建成功）
    if (connId % 100 == 0) {
        std::cout << "Conn " << connId << " added to epoll (fd: " << sockfd << ", port: " << port << ")" << std::endl;
    }
    return sockfd;
}

// 多线程创建连接的工作函数
void EchoClient::createConnectionsThread(int threadId, int totalThreads) {
    for (int i = threadId; i < numConnections_; i += totalThreads) {
        int fd = -1;
        int retry = 0;
        while (retry < 2 && fd == -1) {  // 最多重试1次
            fd = connectToServer(i);
            retry++;
            if (fd == -1 && retry < 2) {
                usleep(50000);  // 重试延迟
            }
        }
        
        // 批量创建后休息，避免系统过载
        if (i % 500 == 0) {
            usleep(100000);
        }
    }
}

void EchoClient::createConnections() {
    std::cout << "=== Starting to create " << numConnections_ << " connections with " << THREAD_NUM << " threads ===" << std::endl;
    std::vector<std::thread> threads;
    
    for (int t = 0; t < THREAD_NUM; ++t) {
        threads.emplace_back(&EchoClient::createConnectionsThread, this, t, THREAD_NUM);
    }
    
    for (auto& th : threads) {
        th.join();  // 等待所有连接创建线程完成
    }
    std::cout << "=== All connection creation threads finished ===" << std::endl;
    std::cout << "Total connections created (including failed retries): " << numConnections_ << std::endl;
}

void EchoClient::eventLoop() {
    struct epoll_event events[MAX_EVENTS];
    char buffer[BUFFER_SIZE];
    std::unordered_map<int, std::pair<std::string, int>> recvBuffer;

    // 总超时设置为60秒（确保有足够时间处理）
    const time_t TOTAL_TIMEOUT = 60;
    time_t startTime = time(nullptr);
    std::cout << "=== Entering event loop (timeout: " << TOTAL_TIMEOUT << "s) ===" << std::endl;

    while (true) {
        // 检查总超时
        time_t currentTime = time(nullptr);
        if (currentTime - startTime > TOTAL_TIMEOUT) {
            std::cerr << "\n=== Event loop timed out after " << TOTAL_TIMEOUT << "s ===" << std::endl;
            break;
        }

        // 等待事件（超时1秒，避免永久阻塞）
        int nfds = epoll_wait(epollFd_, events, MAX_EVENTS, 1000);
        if (nfds == -1) {
            if (errno == EINTR) continue;
            std::cerr << "Epoll wait failed: " << strerror(errno) << std::endl;
            break;
        }

        // 打印事件数量（验证是否有事件）
        if (nfds > 0) {
            std::cout << "Received " << nfds << " events in epoll_wait" << std::endl;
        }

        // 处理超时连接（加锁保护）
        {
            std::lock_guard<std::mutex> lock(connMutex_);
            time_t now = time(nullptr);
            auto it = connCreateTime_.begin();
            while (it != connCreateTime_.end()) {
                int fd = it->first;
                time_t createTime = it->second;
                
                if (now - createTime > CONN_TIMEOUT_SEC) {
                    std::cerr << "Conn fd " << fd << " timed out (>" << CONN_TIMEOUT_SEC << "s)" << std::endl;
                    epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                    close(fd);
                    failedConnections_++;
                    connState_.erase(fd);
                    recvBuffer.erase(fd);
                    it = connCreateTime_.erase(it);  // 安全删除
                } else {
                    ++it;
                }
            }
        }

        // 处理事件
        for (int i = 0; i < nfds; ++i) {
            int fd = events[i].data.fd;
            uint32_t event = events[i].events;

            // 检查连接是否已被清理
            bool isActive = false;
            {
                std::lock_guard<std::mutex> lock(connMutex_);
                isActive = (connCreateTime_.find(fd) != connCreateTime_.end());
            }
            if (!isActive) continue;

            // 处理错误事件
            if (event & EPOLLERR) {
                std::cerr << "Conn fd " << fd << " error event" << std::endl;
                std::lock_guard<std::mutex> lock(connMutex_);
                epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                close(fd);
                failedConnections_++;
                connState_.erase(fd);
                connCreateTime_.erase(fd);
                recvBuffer.erase(fd);
                continue;
            }

            // 处理可写事件（连接完成或发送数据）
            if (event & EPOLLOUT) {
                std::lock_guard<std::mutex> lock(connMutex_);
                auto stateIt = connState_.find(fd);
                if (stateIt == connState_.end()) continue;

                int& echoCount = stateIt->second.second;

                // 首次连接：验证连接是否成功
                if (echoCount == 0) {
                    if (!checkConnectSuccess(fd)) {
                        std::cerr << "Conn fd " << fd << " connect failed" << std::endl;
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        failedConnections_++;
                        connState_.erase(fd);
                        connCreateTime_.erase(fd);
                        continue;
                    }

                    connected_++;
                    std::cout << "Conn fd " << fd << " connected (" << connected_ << "/" << numConnections_ << ")" << std::endl;
                }

                // 生成并发送数据
                std::string testData = generateTestData(fd % numConnections_, echoCount + 1);
                stateIt->second.first = testData;
                int dataLen = testData.size();

                int nwrite = write(fd, testData.c_str(), dataLen);
                if (nwrite == -1) {
                    if (errno != EAGAIN && errno != EWOULDBLOCK) {
                        std::cerr << "Conn fd " << fd << " send failed" << std::endl;
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        failedConnections_++;
                        connState_.erase(fd);
                        connCreateTime_.erase(fd);
                    }
                    continue;
                }

                std::cout << "Conn fd " << fd << " sent (echo " << echoCount + 1 << "): '" << testData << "'" << std::endl;

                // 切换为监听读事件
                struct epoll_event ev;
                ev.events = EPOLLIN | EPOLLERR;
                ev.data.fd = fd;
                epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
            }

            // 处理可读事件（接收回显）
            if (event & EPOLLIN) {
                std::lock_guard<std::mutex> lock(connMutex_);
                auto stateIt = connState_.find(fd);
                if (stateIt == connState_.end()) continue;

                std::string& expectedData = stateIt->second.first;
                int& echoCount = stateIt->second.second;
                int expectedLen = expectedData.size();

                if (recvBuffer.find(fd) == recvBuffer.end()) {
                    recvBuffer[fd] = std::make_pair("", 0);
                }
                std::string& data = recvBuffer[fd].first;
                int& received = recvBuffer[fd].second;

                // 接收数据
                int nbytes = recvAll(fd, buffer, expectedLen - received);
                if (nbytes == -1) {
                    std::cerr << "Conn fd " << fd << " receive failed" << std::endl;
                    epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                    close(fd);
                    failedConnections_++;
                    connState_.erase(fd);
                    connCreateTime_.erase(fd);
                    recvBuffer.erase(fd);
                    continue;
                }

                received += nbytes;
                data.append(buffer, nbytes);
                std::cout << "Conn fd " << fd << " received: '" << data << "' (" << received << "/" << expectedLen << ")" << std::endl;

                // 验证回显
                if (received == expectedLen) {
                    if (data == expectedData) {
                        echoCount++;
                        std::cout << "Conn fd " << fd << " echo success (" << echoCount << "/" << maxEchoCount_ << ")" << std::endl;

                        if (echoCount >= maxEchoCount_) {
                            std::cout << "Conn fd " << fd << " completed all echoes" << std::endl;
                            epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                            close(fd);
                            connState_.erase(fd);
                            connCreateTime_.erase(fd);
                            completed_++;
                            std::cout << "Completed connections: " << completed_ << "/" << numConnections_ << std::endl;
                        } else {
                            // 继续发送下一次数据
                            struct epoll_event ev;
                            ev.events = EPOLLOUT | EPOLLERR;
                            ev.data.fd = fd;
                            epoll_ctl(epollFd_, EPOLL_CTL_MOD, fd, &ev);
                        }
                    } else {
                        std::cerr << "Conn fd " << fd << " mismatch: expected '" << expectedData << "', got '" << data << "'" << std::endl;
                        epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                        close(fd);
                        failedConnections_++;
                        connState_.erase(fd);
                        connCreateTime_.erase(fd);
                    }
                    recvBuffer.erase(fd);
                } else if (received > expectedLen) {
                    std::cerr << "Conn fd " << fd << " invalid length" << std::endl;
                    epoll_ctl(epollFd_, EPOLL_CTL_DEL, fd, nullptr);
                    close(fd);
                    failedConnections_++;
                    connState_.erase(fd);
                    connCreateTime_.erase(fd);
                    recvBuffer.erase(fd);
                }
            }
        }

        // 退出条件：所有连接处理完成
        if (completed_ + failedConnections_ >= numConnections_) {
            std::cout << "\n=== All connections processed ===" << std::endl;
            break;
        }
    }

    // 最终统计
    std::cout << "\n===== Test Result =====" << std::endl;
    std::cout << "Total connections: " << numConnections_ << std::endl;
    std::cout << "Successfully connected: " << connected_ << std::endl;
    std::cout << "Completed all echoes: " << completed_ << std::endl;
    std::cout << "Failed connections: " << failedConnections_ << std::endl;
}

void EchoClient::start() {
    // 先创建所有连接，再进入事件循环（确保连接已注册到epoll）
    std::thread connThread(&EchoClient::createConnections, this);
    connThread.join();  // 等待连接创建完成
    eventLoop();        // 连接创建完成后再处理事件
}